/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "model_handler.h"
#include "utils_log.h"
#include "utils_string_ex.h"
#include "utils_crypto.h"
#include <chrono>
#include <iomanip>
#include <sstream>
#include <fstream>
#include <filesystem>
#include <openssl/md5.h>

namespace El {
namespace Protocol {

void ModelHandler::RegisterMethodHandlers(
    std::unordered_map<std::string, std::function<nlohmann::json(const std::string&, const nlohmann::json&)>>& handlers)
{
    // 注册模型管理相关方法
    handlers["model.upload"] = [this](const std::string& id, const nlohmann::json& params) {
        return HandleModelUpload(id, params);
    };

    handlers["model.list"] = [this](const std::string& id, const nlohmann::json& params) {
        return HandleModelList(id, params);
    };

    handlers["model.download"] = [this](const std::string& id, const nlohmann::json& params) {
        return HandleModelDownload(id, params);
    };

    handlers["model.delete"] = [this](const std::string& id, const nlohmann::json& params) {
        return HandleModelDelete(id, params);
    };
}

nlohmann::json ModelHandler::HandleModelUpload(const std::string& id, const nlohmann::json& params)
{
    // 检查必要参数
    if (!params.contains("chunk_index") || !params["chunk_index"].is_number() || 
        !params.contains("total_chunks") || !params["total_chunks"].is_number() || 
        !params.contains("chunk_data") || !params["chunk_data"].is_string()) {
        return CreateErrorResponse(id, 400, "缺少必要参数或参数格式不正确");
    }

    int chunkIndex = params["chunk_index"];
    int totalChunks = params["total_chunks"];
    std::string chunkData = params["chunk_data"];

    // 第一个分片需要包含文件信息
    if (chunkIndex == 0) {
        if (!params.contains("file_info") || !params["file_info"].is_object()) {
            return CreateErrorResponse(id, 400, "缺少文件信息");
        }

        const auto& fileInfo = params["file_info"];
        if (!fileInfo.contains("filename") || !fileInfo.contains("size")) {
            return CreateErrorResponse(id, 400, "文件信息不完整");
        }

        // 初始化上传状态
        uploadStatus_.fileName = fileInfo.value("filename", "");
        uploadStatus_.totalSize = fileInfo.value("size", 0);
        uploadStatus_.modelDescription = fileInfo.value("description", "");
        uploadStatus_.totalChunks = totalChunks;
        uploadStatus_.uploadedChunks = 0;

        // 创建临时文件路径
        std::string tempDir = "/tmp/models";
        if (!std::filesystem::exists(tempDir)) {
            std::filesystem::create_directories(tempDir);
        }
        
        uploadStatus_.tempFilePath = tempDir + "/" + uploadStatus_.fileName + ".tmp";
        
        // 创建空文件
        std::ofstream file(uploadStatus_.tempFilePath, std::ios::binary);
        if (!file) {
            return CreateErrorResponse(id, 500, "无法创建临时文件");
        }
        file.close();
    }

    // 验证分片序号
    if (chunkIndex >= totalChunks || chunkIndex != uploadStatus_.uploadedChunks) {
        return CreateErrorResponse(id, 400, "分片序号错误");
    }

    // 解码分片数据并写入文件
    std::vector<uint8_t> binaryData = El::Utils::Base64Decode(chunkData);
    
    std::fstream file(uploadStatus_.tempFilePath, std::ios::in | std::ios::out | std::ios::binary);
    if (!file) {
        return CreateErrorResponse(id, 500, "无法打开临时文件");
    }
    
    // 定位到分片位置
    file.seekp(chunkIndex * (1024 * 1024)); // 假设每块1MB
    if (!file) {
        return CreateErrorResponse(id, 500, "无法定位文件位置");
    }
    
    // 写入数据
    file.write(reinterpret_cast<const char*>(binaryData.data()), binaryData.size());
    if (!file) {
        return CreateErrorResponse(id, 500, "写入文件失败");
    }
    
    file.close();
    
    uploadStatus_.uploadedChunks++;

    // 如果是最后一个分片，移动文件到最终位置
    if (uploadStatus_.uploadedChunks == uploadStatus_.totalChunks) {
        std::string modelDir = "/mnt/data/model";
        // 如果无法访问 /mnt/data/model，则使用用户主目录下的替代路径
        if (!std::filesystem::exists("/mnt/data") || !std::filesystem::is_directory("/mnt/data")) {
            const char* homeDir = std::getenv("HOME");
            if (homeDir != nullptr) {
                modelDir = std::string(homeDir) + "/models";
                EL_WARN("/mnt/data 不可用，使用替代路径: {}", modelDir);
            }
        }
        
        if (!std::filesystem::exists(modelDir)) {
            std::filesystem::create_directories(modelDir);
        }
        
        std::string finalPath = modelDir + "/" + uploadStatus_.fileName;
        
        // 复制临时文件到最终位置并删除临时文件
        std::error_code ec;
        EL_INFO("Copying file from {} to {}", uploadStatus_.tempFilePath, finalPath);
        std::filesystem::copy_file(uploadStatus_.tempFilePath, finalPath, ec);
        if (ec) {
            EL_ERROR("Failed to copy file: {}", ec.message());
            return CreateErrorResponse(id, 500, "保存文件失败: " + ec.message());
        }
        
        EL_INFO("File copied successfully, removing temporary file");
        // 删除临时文件
        std::filesystem::remove(uploadStatus_.tempFilePath, ec);
        if (ec) {
            // 即使删除临时文件失败，也不应该影响整体操作
            EL_WARN("Failed to remove temporary file: {}", ec.message());
        }
        EL_INFO("Temporary file removed successfully");
        
        // 保存模型信息到数据库或配置文件
        // 这里简化处理，实际项目中应该保存到数据库
        nlohmann::json modelInfo = {
            {"filename", uploadStatus_.fileName},
            {"size", uploadStatus_.totalSize},
            {"path", finalPath},
            {"description", uploadStatus_.modelDescription},
            {"upload_time", std::time(nullptr)}
        };
        
        // 重置上传状态
        uploadStatus_ = {};
    }

    nlohmann::json result = {{"status", "success"}, {"uploaded_chunks", uploadStatus_.uploadedChunks}};
    return CreateSuccessResponse(id, result);
}

nlohmann::json ModelHandler::HandleModelList(const std::string& id, const nlohmann::json& params)
{
    (void)params; // 未使用参数
    
    // 获取模型列表
    nlohmann::json models = nlohmann::json::array();
    
    std::string modelDir = "/mnt/data/model";
    // 如果无法访问 /mnt/data/model，则使用用户主目录下的替代路径
    if (!std::filesystem::exists("/mnt/data") || !std::filesystem::is_directory("/mnt/data")) {
        const char* homeDir = std::getenv("HOME");
        if (homeDir != nullptr) {
            modelDir = std::string(homeDir) + "/models";
            EL_WARN("/mnt/data 不可用，使用替代路径: {}", modelDir);
        }
    }
    
    if (std::filesystem::exists(modelDir)) {
        for (const auto& entry : std::filesystem::directory_iterator(modelDir)) {
            if (entry.is_regular_file()) {
                nlohmann::json modelInfo = {
                    {"filename", entry.path().filename().string()},
                    {"size", entry.file_size()},
                    {"upload_time", std::filesystem::last_write_time(entry).time_since_epoch().count()}
                };
                models.push_back(modelInfo);
            }
        }
    }

    nlohmann::json result = {{"models", models}};
    return CreateSuccessResponse(id, result);
}

nlohmann::json ModelHandler::HandleModelDownload(const std::string& id, const nlohmann::json& params)
{
    // 检查必要参数
    if (!params.contains("filename") || !params["filename"].is_string() ||
        !params.contains("chunk_index") || !params["chunk_index"].is_number()) {
        return CreateErrorResponse(id, 400, "缺少必要参数");
    }

    std::string filename = params["filename"];
    int chunkIndex = params["chunk_index"];
    int chunkSize = params.value("chunk_size", 1024 * 1024); // 默认1MB
    
    // 构建模型路径，优先使用 /mnt/data/model，如果不可用则使用替代路径
    std::string baseDir = "/mnt/data/model";
    if (!std::filesystem::exists("/mnt/data") || !std::filesystem::is_directory("/mnt/data")) {
        const char* homeDir = std::getenv("HOME");
        if (homeDir != nullptr) {
            baseDir = std::string(homeDir) + "/models";
            EL_WARN("/mnt/data 不可用，使用替代路径: {}", baseDir);
        }
    }
    
    std::string modelPath = baseDir + "/" + filename;

    // 检查文件是否存在
    if (!std::filesystem::exists(modelPath)) {
        return CreateErrorResponse(id, 404, "模型文件不存在");
    }

    // 获取文件大小
    std::uintmax_t fileSize = std::filesystem::file_size(modelPath);
    int totalChunks = (fileSize + chunkSize - 1) / chunkSize;
    
    // 验证分块索引
    if (chunkIndex >= totalChunks) {
        return CreateErrorResponse(id, 400, "分块索引超出范围");
    }

    // 读取指定分块的数据
    std::ifstream file(modelPath, std::ios::binary);
    if (!file) {
        return CreateErrorResponse(id, 500, "无法打开模型文件");
    }
    
    // 定位到分块位置
    file.seekg(chunkIndex * chunkSize);
    if (!file) {
        return CreateErrorResponse(id, 500, "文件定位失败");
    }
    
    // 计算当前分块的实际大小
    int actualChunkSize = std::min(chunkSize, static_cast<int>(fileSize - chunkIndex * chunkSize));
    
    // 读取数据
    std::vector<uint8_t> buffer(actualChunkSize);
    file.read(reinterpret_cast<char*>(buffer.data()), actualChunkSize);
    if (!file && !file.eof()) {
        return CreateErrorResponse(id, 500, "读取文件失败");
    }
    
    // 编码为Base64
    std::string chunkData = El::Utils::Base64Encode(buffer);
    
    nlohmann::json result = {
        {"chunk_data", chunkData},
        {"chunk_index", chunkIndex}
    };
    
    // 如果是第一个分块，添加文件信息
    if (chunkIndex == 0) {
        std::string extension = std::filesystem::path(filename).extension().string();
        std::string format = "unknown";
        if (extension == ".onnx") format = "onnx";
        else if (extension == ".bin") format = "bin";
        else if (extension == ".pt") format = "pt";
        else if (extension == ".om") format = "om";
        
        result["file_info"] = {
            {"filename", filename},
            {"size", fileSize},
            {"total_chunks", totalChunks},
            {"format", format},
            {"checksum", ""} // TODO: 实现文件校验和计算
        };
    }
    
    return CreateSuccessResponse(id, result);
}

nlohmann::json ModelHandler::HandleModelDelete(const std::string& id, const nlohmann::json& params)
{
    // 检查必要参数
    if (!params.contains("filename") || !params["filename"].is_string()) {
        return CreateErrorResponse(id, 400, "缺少模型文件名参数");
    }

    std::string filename = params["filename"];
    
    // 构建模型路径，优先使用 /mnt/data/model，如果不可用则使用替代路径
    std::string baseDir = "/mnt/data/model";
    if (!std::filesystem::exists("/mnt/data") || !std::filesystem::is_directory("/mnt/data")) {
        const char* homeDir = std::getenv("HOME");
        if (homeDir != nullptr) {
            baseDir = std::string(homeDir) + "/models";
            EL_WARN("/mnt/data 不可用，使用替代路径: {}", baseDir);
        }
    }
    
    std::string modelPath = baseDir + "/" + filename;

    // 检查文件是否存在
    if (!std::filesystem::exists(modelPath)) {
        return CreateErrorResponse(id, 404, "模型文件不存在");
    }

    // 删除文件
    std::error_code ec;
    std::filesystem::remove(modelPath, ec);
    if (ec) {
        return CreateErrorResponse(id, 500, "删除文件失败: " + ec.message());
    }

    nlohmann::json result = {{"status", "success"}};
    return CreateSuccessResponse(id, result);
}

} // namespace Protocol
} // namespace El